"""
GStreamer + Enhanced Number ROI YOLO pipeline (0812)

Integrates RTSP frames from `GStreamerWrapper` with the enhanced
`VideoStreamPipelineNumberEnhanced` (number ROI height doubled) for
real-time detection and visualization.

Controls:
- q: quit
- s: save current displayed frame
- d: toggle ROI debug overlay
"""

import os
import sys
import time
from typing import Optional

import cv2

# Import the enhanced video pipeline (doubles number ROI height)
from video_stream_number_enhanced import VideoStreamPipelineNumberEnhanced


def _ensure_rs2_on_path() -> None:
    """Ensure `yolo-jetson/rs2` is on sys.path to import GStreamerWrapper."""
    current_dir = os.path.dirname(os.path.abspath(__file__))
    rs2_dir = os.path.normpath(os.path.join(current_dir, '..', 'rs2'))
    if rs2_dir not in sys.path:
        sys.path.append(rs2_dir)


def _find_model_weights() -> Optional[str]:
    """Find available model weights using common candidate locations."""
    # Allow override via environment variable
    env_model = os.environ.get('YOLO_MODEL_PATH')
    if env_model and os.path.exists(env_model):
        return env_model

    candidate_models = [
        r".\weights\all_yolo_all_epochs_0811\best.pt",
        r".\weights\all_yolo_all_epochs_0808\best.pt",
        r".\weights\all_yolo_all_epochs_0805\best.pt",
        r".\weights\all_yolo_all_epochs_0804\best.pt",
    ]
    for path in candidate_models:
        if os.path.exists(path):
            return path
    return None


def main() -> int:
    # Resolve model
    model_path = _find_model_weights()
    if model_path is None:
        print("❌ Model weights not found. Set YOLO_MODEL_PATH env var or place weights under .\\weights\\...")
        return 1

    # Resolve and import GStreamer wrapper
    _ensure_rs2_on_path()
    try:
        from GStreamerWrapper import GStreamerWrapper  # type: ignore
    except Exception as e:
        print(f"❌ Failed to import GStreamerWrapper: {e}")
        return 1

    # Initialize GStreamer RTSP source
    try:
        gs = GStreamerWrapper(width=1920, height=1080)
        print("📡 GStreamer RTSP pipeline initialized")
    except Exception as e:
        print(f"❌ Unable to initialize GStreamer pipeline: {e}")
        return 1

    # Create enhanced detection/display pipeline
    print("🔧 Creating NUMBER ROI ENHANCED video stream pipeline")
    print(f"📦 Model: {model_path}")
    try:
        video_pipeline = VideoStreamPipelineNumberEnhanced(
            model_path,
            use_fresh_pipeline_per_frame=True,
        )
    except Exception as e:
        print(f"❌ Failed to initialize VideoStreamPipelineNumberEnhanced: {e}")
        try:
            gs.StopThread()
        except Exception:
            pass
        return 1

    # Runtime variables
    frame_count = 0
    successful_detections = 0
    roi_debug = True
    out_writer: Optional[cv2.VideoWriter] = None
    current_fps = 0.0
    fps_counter = 0
    fps_timer = cv2.getTickCount()

    roi_stats = {
        'total_rois': 0,
        'number_rois': 0,
        'cone_rois': 0,
        'dashboard_rois': 0,
    }

    print("🚀 Starting RTSP stream processing (NUMBER ROI ENHANCED). Press 'q' to quit, 's' to save frame, 'd' to toggle ROI debug mode.")

    try:
        while True:
            frame = gs.GetFrame()
            if frame is None:
                time.sleep(0.01)
                continue

            try:
                results = video_pipeline.process_single_frame_enhanced(frame, frame_count)

                cone_color = results.get('cone_color', 'None')
                number = results.get('number', 'None')
                dashboard_state = results.get('dashboard_state', 'None')

                has_detections = (
                    cone_color != 'None' or number != 'None' or dashboard_state != 'None'
                )
                if has_detections:
                    successful_detections += 1

                # ROI stats
                info = results.get('processing_info', {})
                roi_types = info.get('roi_types', []) if isinstance(info, dict) else []
                roi_stats['total_rois'] += len(roi_types)
                roi_stats['number_rois'] += roi_types.count('number')
                roi_stats['cone_rois'] += roi_types.count('cone')
                roi_stats['dashboard_rois'] += roi_types.count('dashboard')

                # Display
                success_rate = (successful_detections / (frame_count + 1)) * 100 if frame_count >= 0 else 0.0
                display_frame = video_pipeline.create_enhanced_display(
                    frame, results, frame_count, current_fps, success_rate
                )

                if roi_debug and frame_count > 0:
                    debug_y = display_frame.shape[0] - 120
                    avg_rois = roi_stats['total_rois'] / (frame_count + 1)
                    roi_debug_lines = [
                        f"Avg ROIs/frame: {avg_rois:.1f}",
                        f"Number ROIs: {roi_stats['number_rois']} (Enhanced Height)",
                        f"Cone ROIs: {roi_stats['cone_rois']} (Width Optimized)",
                        f"Dashboard ROIs: {roi_stats['dashboard_rois']}",
                    ]
                    for line in roi_debug_lines:
                        video_pipeline._draw_text_with_background(
                            display_frame, line, (30, debug_y),
                            font_color=video_pipeline.colors.get('enhanced', video_pipeline.colors['warning']),
                            font_scale=0.7, thickness=1,
                        )
                        debug_y += 25

                # FPS update
                fps_counter += 1
                if fps_counter % 10 == 0:
                    current_time = cv2.getTickCount()
                    elapsed = (current_time - fps_timer) / cv2.getTickFrequency()
                    if elapsed > 0:
                        current_fps = fps_counter / elapsed
                    fps_counter = 0
                    fps_timer = current_time

                cv2.imshow('YOLO Traffic Light Detection - NUMBER ROI ENHANCED (RTSP)', display_frame)
                if out_writer is not None:
                    out_writer.write(display_frame)

                key = cv2.waitKey(1) & 0xFF
                if key == ord('q'):
                    print("👋 User requested exit")
                    break
                elif key == ord('s'):
                    save_path = f"enhanced_rtsp_frame_{frame_count:06d}.jpg"
                    cv2.imwrite(save_path, display_frame)
                    print(f"💾 Frame saved: {save_path}")
                elif key == ord('d'):
                    roi_debug = not roi_debug
                    print(f"🔧 ROI debug mode: {'enabled' if roi_debug else 'disabled'}")

                frame_count += 1

            except Exception as e:
                print(f"❌ Error processing frame {frame_count}: {e}")
                # continue loop

    except KeyboardInterrupt:
        print("⚠️ User interrupted")
    finally:
        try:
            gs.StopThread()
        except Exception:
            pass
        del gs
        cv2.destroyAllWindows()

        # Final stats
        print("\n🏁 RTSP processing completed (NUMBER ROI ENHANCED)")
        print(f"📊 Total frames processed: {frame_count}")
        print(f"✅ Successful detections: {successful_detections}")
        if frame_count > 0:
            final_success_rate = (successful_detections / frame_count) * 100
            print(f"📈 Overall success rate: {final_success_rate:.1f}%")
            if roi_stats['total_rois'] > 0:
                print("\n🔧 ROI generation statistics (Enhanced):")
                print(f"   Total ROIs: {roi_stats['total_rois']}")
                print(f"   Average per frame: {roi_stats['total_rois'] / frame_count:.1f}")
                print(f"   Number ROIs: {roi_stats['number_rois']} ({roi_stats['number_rois']/roi_stats['total_rois']*100:.1f}%) - Enhanced Height")
                print(f"   Cone ROIs: {roi_stats['cone_rois']} ({roi_stats['cone_rois']/roi_stats['total_rois']*100:.1f}%) - Width Optimized")
                print(f"   Dashboard ROIs: {roi_stats['dashboard_rois']} ({roi_stats['dashboard_rois']/roi_stats['total_rois']*100:.1f}%)")

    return 0


if __name__ == "__main__":
    raise SystemExit(main())


